package com.hitqz.robot.biz.websocket.endpoint;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.collection.ConcurrentHashSet;
import cn.hutool.core.net.NetUtil;
import cn.hutool.core.util.HashUtil;
import cn.hutool.crypto.digest.DigestUtil;
import cn.hutool.crypto.digest.MD5;
import cn.hutool.extra.spring.SpringUtil;
import cn.hutool.system.SystemUtil;
import com.hitqz.robot.biz.controller.RtspController;
import jakarta.websocket.*;
import jakarta.websocket.server.ServerEndpoint;
import lombok.AllArgsConstructor;
import lombok.SneakyThrows;
import lombok.experimental.Accessors;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.web.ServerProperties;
import org.springframework.stereotype.Component;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.net.InetAddress;
import java.nio.ByteBuffer;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

@ServerEndpoint(value = "/live")
@Component
@Slf4j
public class LiveEndpoint {
    // 存储sourceKey和会话的对应关系
    private static final Map<String, Set<Session>> sourceKeySessionMap = new ConcurrentHashMap<>();
    // 用于存储所有打开的WebSocket sessions
    private static final Set<Session> sessions = new CopyOnWriteArraySet<>();
    // 用于存储视频源地址
    private static final Set<String> sourcesKeySet = new ConcurrentHashSet<>();


    @OnOpen
    public void onOpen(Session session) throws IOException {
        sessions.add(session);
        log.debug("[ws][live]session-id:{} connected.current size :{}", session.getId(), sessions.size());
        String queryString = session.getQueryString();
        String[] pairs = queryString.split("&");
        Map<String, String> queryParams = new HashMap<>();
        for (String pair : pairs) {
            // 使用=分割键和值
            int idx = pair.indexOf("=");
            // 将键值对添加到Map中
            queryParams.put(pair.substring(0, idx), pair.substring(idx + 1));
        }
        String source = queryParams.get("source");
        String rs = Optional.ofNullable(queryParams.get("rs")).orElse("640X480");
        String sourceKey = DigestUtil.md5Hex(source + rs);
        //保存 sourceKey和session的对应关系 一对多
        if (sourceKeySessionMap.getOrDefault(sourceKey, Collections.emptySet()).contains(session)) {
            session.close(new CloseReason(CloseReason.CloseCodes.VIOLATED_POLICY, "Duplicate connection"));
            log.warn("[ws][live]sourceKey:{},session-id:{},duplicate connection", sourceKey, session.getId());
        }
        sourceKeySessionMap.computeIfAbsent(sourceKey, k -> ConcurrentHashMap.newKeySet()).add(session);
        synchronized (sourceKey.intern()) {
            if (sourcesKeySet.add(sourceKey)) {
                // 开始推流
                // 开启视频转码线程
                ExecutorService videoPushThreadPool = SpringUtil.getBean("videoPushThreadPool");
                ServerProperties serverProperties = SpringUtil.getBean(ServerProperties.class);
                log.info("[ws][live]sourceKey:{},source:{}，start push", sourceKey, source);
                videoPushThreadPool.submit(new VideoTranscodeThread(source, rs, videoPushThreadPool, serverProperties));
            }
        }
    }

    @OnMessage
    public void onMessage(Session session, String message) {
        log.info("[ws][live]receive session-id:{},message:{}", session.getId(), message);
    }

    @OnClose
    public void onClose(Session session) {
        // 移除session
        sessions.remove(session);
        log.info("[ws][live]session-id:{},disconnected.current size :{}", session.getId(), sessions.size());
        for (Set<Session> sessions : sourceKeySessionMap.values()) {
            sessions.remove(session);
        }
        sourceKeySessionMap.entrySet().removeIf(entry -> entry.getValue().isEmpty());
    }

    @OnError
    public void onError(Session session, Throwable error) {
        log.error("[ws][live]session-id:{},error:{}", session.getId(), error.getMessage());
    }

    public static void sendVideoData(String sourceKey, ByteBuffer buffer) {
        Set<Session> sessions = sourceKeySessionMap.get(sourceKey);
        if (CollUtil.isEmpty(sessions)) {
            return;
        }
        // 创建一个线程池，用于处理不同 session 的并行任务
        ExecutorService websocketSendThreadPool = SpringUtil.getBean("websocketSendThreadPool");
        // 创建一个 CountDownLatch，计数器等于 session 的数量
        CountDownLatch latch = new CountDownLatch(sessions.size());
        // 为每个 session 提交异步任务
        for (Session session : sessions) {
            if (!session.isOpen()) {
                latch.countDown();  // 如果 session 已关闭，减少计数
                continue;
            }
            // 提交异步任务到线程池
            websocketSendThreadPool.submit(() -> {
                try {
                    synchronized (session) {  // 同一个 session 同步发送
                        ByteBuffer sessionBuffer = buffer.duplicate();  // 创建独立的 ByteBuffer
                        session.getBasicRemote().sendBinary(sessionBuffer);  // 发送数据
                    }
                } catch (Exception e) {
                    log.error("[ws][live]send binary data error", e);
                } finally {
                    latch.countDown();  // 任务完成后，减少计数
                }
            });
        }
        // 等待所有 session 发送任务完成
        try {
            latch.await();  // 阻塞直到所有 session 完成
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();  // 保持中断状态
            log.error("Video data sending was interrupted", e);
        }
    }





    // 视频转码线程类
    public static class VideoTranscodeThread implements Runnable {

        private final String source;
        private final String rs;
        private final AtomicLong pid = new AtomicLong(-1);

        private final ExecutorService videoPushThreadPool;
        private final ServerProperties serverProperties;

        public VideoTranscodeThread(String source, String rs, ExecutorService videoPushThreadPool, ServerProperties serverProperties) {
            this.source = source;
            this.rs = rs;
            this.videoPushThreadPool = videoPushThreadPool;
            this.serverProperties = serverProperties;
        }


        @SneakyThrows
        @Override
        public void run() {
            String ip = source.split("@")[1].split(":554")[0];
            String channel = source.split("/")[source.split("/").length - 1];
            while (true) {
                try {
                    while (!NetUtil.ping(ip, 5000)) {
                        log.info("[pushThread] ip:{},channel:{} ping false", ip, channel);
                    }
                    log.info("[pushThread] ip:{},channel:{} ping true", ip, channel);
                    // 推流
                    transcodeAndStream(source, rs);
                    Thread.sleep(200);
                } catch (Exception e) {
                    if (pid.get() != -1) {
                        ProcessHandle.of(pid.get()).ifPresent(ProcessHandle::destroyForcibly);  // 强制销毁已有进程
                        pid.set(-1);
                    }
                    log.error("[pushThread] transcodeAndStream error {},and restart...", e.getMessage());
                    Thread.sleep(3000);  // 3秒后重试
                }
            }
        }

        public void transcodeAndStream(String source, String rs) throws IOException, InterruptedException {
            String sourceKey = DigestUtil.md5Hex(source + rs);
            String output = "http://127.0.0.1:" + serverProperties.getPort() + "/rtsp/ws?sourceKey=" + sourceKey;
            String command = "ffmpeg -i " + source + " -c:v mpeg1video -f mpegts -s " + rs + " " + output;

            ProcessBuilder processBuilder = new ProcessBuilder(command.split(" "));
            Process process = processBuilder.start();

            AtomicBoolean outputAlive = new AtomicBoolean(true);
            // 存储进程 ID
            pid.set(process.pid());
            log.info("[pushThread] ffmpeg process pid:{},command:{}", pid, command);
            // 处理标准输出流
            videoPushThreadPool.submit(() -> {
                try (BufferedReader reader = new BufferedReader(
                        new InputStreamReader(process.getInputStream()))) {
                    String line;
                    while ((line = reader.readLine()) != null) {
                        // log.debug("[pushThread] ffmpeg output: {}" , line);
                        outputAlive.set(true);  // 重置标志，表示有输出数据
                    }
                } catch (IOException e) {
                    log.error("[pushThread] ffmpeg input stream log thread error", e);
                }
            });
            // 处理错误输出流
            videoPushThreadPool.submit(() -> {
                try (BufferedReader reader = new BufferedReader(
                        new InputStreamReader(process.getErrorStream()))) {
                    String line;
                    while ((line = reader.readLine()) != null) {
                        //log.error("[pushThread] ffmpeg error: {}", line);
                        outputAlive.set(true);  // 重置标志，表示有输出数据
                    }
                } catch (IOException e) {
                    log.error("[pushThread] ffmpeg error stream log thread error", e);
                }
            });
            videoPushThreadPool.submit(() -> {
                while (true) {
                    try {
                        Thread.sleep(60*1000);  // 每10秒检测一次
                        if (!outputAlive.getAndSet(false)) {
                            log.warn("No ffmpeg output, restarting stream...");
                            if (pid.get() != -1) {
                                ProcessHandle.of(pid.get()).ifPresent(ProcessHandle::destroyForcibly);  // 强制销毁已有进程
                                pid.set(-1);
                            }
                            break;
                        }
                    } catch (Exception e) {
                        log.error("[pushThread] ffmpeg check output thread error", e);
                    }
                }
            });
            // 等待 ffmpeg 进程结束
            int i = process.waitFor();
            log.warn("[pushThread] ffmpeg exit code :{}", i);
        }
    }


}
